from config.LoadConfig import get_config
from config.clip_classify_config import *
from utils.clip_classify_util import *
from utils import DataUtil
from config.ServerApp import app
from PIL import Image
import requests
import numpy as np
import io
import traceback
import time

_system_conf = get_config("system")


def get_img_prob(img, texts_list, model, processor):
    import time

    clip_start = time.time()
    device = next(model.parameters()).device

    inputs = processor(text=texts_list, images=img, return_tensors="pt", padding=True)
    inputs = inputs.to(device)

    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
    probs = logits_per_image.softmax(dim=1)
    # 格式化为普通小数并保留5位小数
    formatted_probs = [[value for value in row] for row in probs.detach().cpu().numpy()]

    clip_time = time.time() - clip_start
    print(f"[CLIP推理] 单次推理耗时: {clip_time:.3f}秒，设备: {device}，文本数量: {len(texts_list)}")

    return formatted_probs[0]


def single_classify(img_url):
    """
    :param img_url: 图片url
    :return: 分类概率
    """
    # 记录开始时间
    start_time = time.time()
    processed_images = 1  # 单张图片处理

    # 初始化返回字段值
    img_md5 = None
    is_processed = False
    is_gamb = None
    cls_prob = None

    try:
        print(f"开始单张图片CLIP分类")
        app.ctx.logger.info(f"开始单张图片CLIP分类")
        # 加载模型
        model, processor = get_clip_model(CLIP_MODEL_PATH)
        text_list = GAMB_MERGE_LIST

        # 从url加载图片
        response = requests.get(img_url, timeout=10)  # 设置超时
        response.raise_for_status()  # 如果请求失败则抛出HTTPError
        image_bytes = response.content
        img = Image.open(io.BytesIO(image_bytes))
        # 确保图片是RGB格式，有些png可能是RGBA或者P模式
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # 计算MD5
        img_md5 = calculate_md5_from_bytes(image_bytes)

        probs_list = get_img_prob(img, text_list, model, processor)
        cls_prob = probs_list[0] + probs_list[1]

        if cls_prob >= 0.5:
            is_gamb = True
            cls_prob = float(cls_prob)  # 确保是float类型
        elif cls_prob < 0.5:
            is_gamb = False
            cls_prob = float(cls_prob)  # 确保是float类型
        is_processed = True

        # 计算处理时间并输出统计信息
        end_time = time.time()
        processing_time = end_time - start_time
        print(f"单张图片CLIP分类完成！处理图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒")
        app.ctx.logger.info(f"单张图片CLIP分类统计 - 图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒")

        img_data = {
            "img_md5": img_md5,
            "is_processed": is_processed,
            "is_gamb": is_gamb,
            "cls_prob": cls_prob
        }

        return img_data

    except Exception as e:
        # 即使出现异常也记录处理时间
        end_time = time.time()
        processing_time = end_time - start_time
        print(f"单张图片CLIP分类异常！已处理图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        app.ctx.logger.error(f"单张图片CLIP分类异常 - 图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        raise e


def clip_classify(img_path):
    """
    :param img_url: 图片url
    :param clip_type: 分类类型(可填：'gamb', 'custom')
    :return: 分类概率
    """
    result = []
    # 记录开始时间
    start_time = time.time()
    processed_images = 1  # 单张图片处理

    try:
        clip_type = _system_conf.get("clip_type", "gamb")
        print(f"开始CLIP分类，类型: {clip_type}")
        app.ctx.logger.info(f"开始CLIP分类，类型: {clip_type}")
        # 加载模型
        model, processor = get_clip_model(CLIP_MODEL_PATH)

        # 统一加载图片（避免重复加载）
        img = Image.open(img_path)
        if img.mode != 'RGB':
            img = img.convert('RGB')

        # 初始化赌博概率变量（避免作用域问题）
        gamb_related_prob = 0.0

        if "gamb" in clip_type:
            # 初始化返回字段值
            first_text_list = GAMB_MERGE_LIST
            second_text_list = GAMB_SECOND_LIST

            first_probs_list = get_img_prob(img, first_text_list, model, processor)
            first_cls_prob = first_probs_list[0] + first_probs_list[1] + first_probs_list[2]
            normal_prob = first_probs_list[3] if len(first_probs_list) > 3 else 0

            print(f"CLIP赌博分类概率详情: {first_probs_list}")
            print(f"相关概率: {first_cls_prob:.4f}, 正常概率: {normal_prob:.4f}")

            # 更严格的判断：提高阈值 + 相对优势
            if first_cls_prob >= 0.95 and first_cls_prob > normal_prob * 1.5:
                is_gamb = True
                second_probs_list = get_img_prob(img, second_text_list, model, processor)
                second_cls_prob = max(second_probs_list)

                first_cls_name = "赌博"
                second_cls_name = GAMB_SECOND_SHOW_LIST[np.argmax(second_probs_list)]
                second_cls_prob = float(second_cls_prob)  # 确保是float类型

                data_result = DataUtil.build_classfy_third_obj(first_cls_name, second_cls_name, second_cls_name,
                                                               str(np.round(second_cls_prob, 4)))

                result.append(data_result)

        if "eroticism" in clip_type:
            # 初始化返回字段值
            first_text_list = EROTICISM_MERGE_LIST
            second_text_list = EROTICISM_SECOND_LIST

            first_probs_list = get_img_prob(img, first_text_list, model, processor)
            first_cls_prob = first_probs_list[0] + first_probs_list[1] + first_probs_list[2]
            normal_prob = first_probs_list[3] if len(first_probs_list) > 3 else 0

            print(f"CLIP涉黄分类概率详情: {first_probs_list}")
            print(f"相关概率: {first_cls_prob:.4f}, 正常概率: {normal_prob:.4f}")

            # 涉黄需要更严格的判断避免误报
            if first_cls_prob >= 0.95 and first_cls_prob > normal_prob * 2.0:
                second_probs_list = get_img_prob(img, second_text_list, model, processor)
                second_cls_prob = max(second_probs_list)

                first_cls_name = "涉黄"
                second_cls_name = EROTICISM_SECOND_SHOW_LIST[np.argmax(second_probs_list)]
                second_cls_prob = float(second_cls_prob)  # 确保是float类型

                data_result = DataUtil.build_classfy_third_obj(first_cls_name, second_cls_name, second_cls_name,
                                                               str(np.round(second_cls_prob, 4)))

                result.append(data_result)

        if "weapon" in clip_type:
            # 初始化返回字段值
            first_text_list = WEAPON_MERGE_LIST
            second_text_list = WEAPON_SECOND_LIST

            first_probs_list = get_img_prob(img, first_text_list, model, processor)
            first_cls_prob = first_probs_list[0] + first_probs_list[1] + first_probs_list[2]
            normal_prob = first_probs_list[3] if len(first_probs_list) > 3 else 0

            print(f"CLIP枪械刀具分类概率详情: {first_probs_list}")
            print(f"相关概率: {first_cls_prob:.4f}, 正常概率: {normal_prob:.4f}")

            # 枪械刀具严格判断避免误报
            if first_cls_prob >= 0.95 and first_cls_prob > normal_prob * 1.5:
                second_probs_list = get_img_prob(img, second_text_list, model, processor)
                second_cls_prob = max(second_probs_list)

                first_cls_name = "枪械刀具"
                second_cls_name = WEAPON_SECOND_SHOW_LIST[np.argmax(second_probs_list)]
                second_cls_prob = float(second_cls_prob)  # 确保是float类型

                data_result = DataUtil.build_classfy_third_obj(first_cls_name, second_cls_name, second_cls_name,
                                                               str(np.round(second_cls_prob, 4)))

                result.append(data_result)

        if "violence" in clip_type:
            # 初始化返回字段值
            first_text_list = VIOLENCE_MERGE_LIST
            second_text_list = VIOLENCE_SECOND_LIST

            first_probs_list = get_img_prob(img, first_text_list, model, processor)
            first_cls_prob = first_probs_list[0] + first_probs_list[1] + first_probs_list[2]
            normal_prob = first_probs_list[3] if len(first_probs_list) > 3 else 0

            print(f"CLIP血腥暴力分类概率详情: {first_probs_list}")
            print(f"相关概率: {first_cls_prob:.4f}, 正常概率: {normal_prob:.4f}")

            # 血腥暴力严格判断避免误报
            if first_cls_prob >= 0.95 and first_cls_prob > normal_prob * 1.5:
                second_probs_list = get_img_prob(img, second_text_list, model, processor)
                second_cls_prob = max(second_probs_list)

                first_cls_name = "血腥暴力"
                second_cls_name = VIOLENCE_SECOND_SHOW_LIST[np.argmax(second_probs_list)]
                second_cls_prob = float(second_cls_prob)  # 确保是float类型

                data_result = DataUtil.build_classfy_third_obj(first_cls_name, second_cls_name, second_cls_name,
                                                               str(np.round(second_cls_prob, 4)))

                result.append(data_result)

        if "custom" in clip_type:
            # 初始化返回字段值
            first_text_list = CUSTOM_LIST
            second_text_list = CUSTOM_SHOW_LIST

            # 图片已在上方统一加载

            probs_list = get_img_prob(img, first_text_list, model, processor)

            max_prob_index = np.argmax(probs_list)
            max_prob = max(probs_list)
            custom_cls_name = second_text_list[max_prob_index]
            custom_cls_prob = float(max_prob)  # 确保是float类型

            print(f"CLIP海关分类概率: {max_prob}, 索引: {max_prob_index}")
            app.ctx.logger.info(f"CLIP海关分类概率: {max_prob}, 索引: {max_prob_index}")

            # 添加概率阈值，避免误分类
            if max_prob_index != 2 and max_prob > 0.6:
                data_result = DataUtil.build_classfy_third_obj("海关", custom_cls_name, custom_cls_name,
                                                               str(np.round(custom_cls_prob, 4)))
                result.append(data_result)

        # 计算处理时间并输出统计信息
        end_time = time.time()
        processing_time = end_time - start_time
        print(
            f"CLIP分类完成！处理图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒，平均每张: {processing_time / processed_images:.2f}秒")
        app.ctx.logger.info(
            f"CLIP分类统计 - 图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒，平均每张: {processing_time / processed_images:.2f}秒")

        return result
    except Exception as e:
        # 即使出现异常也记录处理时间
        end_time = time.time()
        processing_time = end_time - start_time
        print(f"CLIP分类异常！已处理图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        app.ctx.logger.error(f"CLIP分类异常 - 图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        traceback.print_exc()
        raise e


def gamb_classify_batch(local_img_path):
    """
    :param img_url: 图片url
    :return: 分类概率
    """
    # 记录开始时间
    start_time = time.time()
    processed_images = 1  # 单张图片处理

    # 初始化返回字段值

    try:
        print(f"开始批量CLIP分类")
        app.ctx.logger.info(f"开始批量CLIP分类")
        # 加载模型
        model, processor = get_clip_model(CLIP_MODEL_PATH)

        img = Image.open(local_img_path)
        # 确保图片是RGB格式，有些png可能是RGBA或者P模式
        if img.mode != 'RGB':
            img = img.convert('RGB')
        result = []

        for cls_dict in CLS_BATCH_DICT:
            cls_name = cls_dict.get("cls_name")
            first_text_list = cls_dict.get("cls_list", [])
            score_reduce_count = cls_dict.get("score_reduce_count", 1)
            second_text_list = cls_dict.get("sub_cls_list", [])

            if not first_text_list:
                continue

            first_probs_list = get_img_prob(img, first_text_list, model, processor)
            if score_reduce_count > 2:
                first_cls_prob = sum(first_probs_list[:score_reduce_count])
            else:
                first_cls_prob = first_probs_list[0]

            if first_cls_prob >= 0.5:
                # 从sub_cls_list中获取语义描述列表
                second_text_list_keys = list(second_text_list[0].keys())
                second_probs_list = get_img_prob(img, second_text_list_keys, model, processor)

                second_cls_prob = max(second_probs_list)

                # 获取对应的二级分类简称
                second_cls_name = list(second_text_list[0].values())[np.argmax(second_probs_list)]
                second_cls_prob = float(second_cls_prob)  # 确保是float类型

                data_result = DataUtil.build_classfy_third_obj(cls_name, second_cls_name, second_cls_name,
                                                               str(np.round(second_cls_prob, 4)))
                result.append(data_result)

        # 计算处理时间并输出统计信息
        end_time = time.time()
        processing_time = end_time - start_time
        print(
            f"批量CLIP分类完成！处理图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒，平均每张: {processing_time / processed_images:.2f}秒")
        app.ctx.logger.info(
            f"批量CLIP分类统计 - 图片数量: {processed_images}张，总耗时: {processing_time:.2f}秒，平均每张: {processing_time / processed_images:.2f}秒")

        return result
    except Exception as e:
        # 即使出现异常也记录处理时间
        end_time = time.time()
        processing_time = end_time - start_time
        print(f"批量CLIP分类异常！已处理图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        app.ctx.logger.error(f"批量CLIP分类异常 - 图片数量: {processed_images}张，耗时: {processing_time:.2f}秒")
        traceback.print_exc()
        raise e

